import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载微调后的模型和分词器
model_name = "./my_finetuned_model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def chat_with_model():
    print("Starting chat with AI model. Type 'exit' to end the conversation.")
    history = []
    
    while True:
        user_input = input("\nYou: ")
        if user_input.lower() == 'exit':
            print("Ending conversation.")
            break
        
        # 将用户输入添加到历史记录中
        history.append(tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt'))
        
        # 将历史记录拼接起来
        input_ids = torch.cat(history, dim=-1)
        
        # 生成模型的响应
        response = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
        
        # 提取模型生成的文本
        response_text = tokenizer.decode(response[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
        
        print(f"AI: {response_text}")
        
        # 将模型的响应添加到历史记录中
        history.append(tokenizer.encode(response_text + tokenizer.eos_token, return_tensors='pt'))

if __name__ == "__main__":
    chat_with_model()